import json
import argparse

def convert_json(input_file, output_file):
    with open(input_file, 'r', encoding='utf-8') as f:
        try:
            data = json.load(f)  # Try loading as JSON array
        except json.JSONDecodeError:
            # If not a valid JSON array, try line by line
            f.seek(0)
            data = [json.loads(line) for line in f if line.strip()]

    instances = []
    for item in data:
        answer = item.pop("answer")
        input_data = item
        instances.append({
            "input": json.dumps(input_data, ensure_ascii=False),
            "output": str(answer)
        })

    output = {
        "type": "text2text",
        "instances": instances
    }

    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(output, f, indent=2, ensure_ascii=False)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert JSON dataset to input-output format.")
    parser.add_argument("--input", type=str, required=True, help="Path to input JSON file")
    parser.add_argument("--output", type=str, required=True, help="Path to output JSON file")
    args = parser.parse_args()

    convert_json(args.input, args.output)